# -*- coding: utf-8 -*-
"""
Created on Wed Jun  9 22:38:01 2021

@author: Administrator
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif']=['SimHei']    #画图时使用中文字体
plt.rcParams['axes.unicode_minus'] = False

def plot_plate_index(plate_index, window=False, future=False, shanghai=False):
    '''
    画图
    '''
    fig = plt.figure(figsize=(12, 6))
    plt.plot(plate_index.index, plate_index.values)
    plt.xlabel('日期', fontsize=16)
    plt.ylabel('收盘指数', fontsize=16)
    if window:
        text = f'移动平均{window}的'
    else:
        text = ''
    if shanghai:
        plt.title(text+'上证指数', fontsize=24)
    else:
        plt.title(text+'光伏-建筑板块指数', fontsize=24)
    plt.grid()
    if future:
        fig.savefig(f'../图片/未来{text}板块指数.png')
    else:
        fig.savefig(f'../图片/{text}板块指数.png')
    plt.show()
    

def return_plate_index(path):
    '''
    计算光伏-建筑板块指数
    '''
    capitals = pd.read_excel(path, index_col=0)
    # 返回表中每一行的缺失值数量，意味着当天，带有缺失值的证券，缺失是因为没上市。
    nan_num = capitals.apply(lambda x: sum(x.isna()), axis=1)
    # 增加一列
    capitals['nan'] = nan_num
    nan_num_list = list(set(nan_num))
    nam_num_list = sorted(nan_num_list, reverse=True)
    plate_index = pd.DataFrame()
    for i in nan_num_list:
        sub_capitals = capitals.loc[capitals['nan']==i]
        # 计算基期调整市值
        base = capitals.loc[capitals['nan']==i].iloc[0, :-1]
        sub_capitals.drop(['nan'], inplace=True, axis=1)
        # 证券基期总市值
        base_total = np.sum(base)
        # 当日的总市值
        total = sub_capitals.sum(axis=1)
        # 证券调整市值 X 权重因子
        sub_capitals = sub_capitals.apply(lambda x: x*x/total)
    
        sub_capitals = (sub_capitals/base_total).sum(axis=1) * 1000
        # 
        plate_index = pd.concat([plate_index, sub_capitals], \
                                ignore_index = False)
        
    plate_index.columns = ['光伏-建筑板块指数']
    plate_index = plate_index.sort_index()
    return plate_index

if __name__ == '__main__':
    path = r'../附件/中间数据/证券市值2019-4-1_2021-5-28.xlsx'
    # plate_index = pd.read_excel(r'../附件/中间数据/光伏-建筑板块市值.xlsx',index_col=0)
    plate_index = return_plate_index(path)
    plate_index.to_excel(r'../附件/中间数据/光伏-建筑板块市值.xlsx')
    
    # 画图（原始数据），画出 2019-4-1 到 2021-4-30 的时序数据
    plate_index_analyze = plate_index['2019-4-1':'2021-4-30']
    plot_plate_index(plate_index_analyze)
    for window in [5,10,20]:
        plate_index_MA = plate_index_analyze.rolling(window=window).mean()
        plot_plate_index(plate_index_MA, window=window)